package formal.tree;

/**
 * TrieNode，又称前缀树或字典树，用于判断字符串是否存在或者是否具有某种字符串前缀
 * https://leetcode-cn.com/problems/implement-trie-prefix-tree/solution/shi-xian-trie-qian-zhui-shu-by-leetcode/
 *
 * @author DengYuan2
 * @create 2021-01-26 21:59
 */
public class Trie {
    private TrieNode root;

    public Trie() {
        root = new TrieNode();
    }

    public void insert(String word) {
        TrieNode node = root;
        for (int i = 0; i < word.length(); i++) {
            char c = word.charAt(i);
            if (!node.containsKey(c)) {
                node.put(c, new TrieNode());
            }
            node = node.get(c);
        }
        node.setEnd();
    }

    /**
     * 查找键，要求完全匹配到叶子节点
     *
     * @param word
     * @return
     */
    public boolean search(String word) {
        TrieNode trieNode = searchPrefix(word);
        return trieNode != null && trieNode.isEnd();
    }

    public TrieNode searchPrefix(String word) {
        TrieNode node = root;
        for (int i = 0; i < word.length(); i++) {
            char c = word.charAt(i);
            if (node.containsKey(c)) {
                node = node.get(c);
            } else {
                return null;
            }
        }
        return node;
    }

    public boolean startsWith(String prefix) {
        TrieNode trieNode = searchPrefix(prefix);
        return trieNode != null;
    }

}


class TrieNode {
    /**
     * 孩子节点
     */
    private TrieNode[] children;
    private final int amount = 26;
    /**
     * 该节点是否是叶子节点
     */
    private boolean isEnd;

    public TrieNode() {
        children = new TrieNode[amount];
    }

    /**
     * 是否有某孩子节点
     *
     * @param ch
     * @return
     */
    public boolean containsKey(char ch) {
        return children[ch - 'a'] != null;
    }

    /**
     * 得到某孩子节点
     *
     * @param ch
     * @return
     */
    public TrieNode get(char ch) {
        return children[ch - 'a'];
    }

    /**
     * 将节点放到合适的位置
     *
     * @param ch
     * @param node
     */
    public void put(char ch, TrieNode node) {
        children[ch - 'a'] = node;
    }

    /**
     * 设置该节点为叶子节点
     */
    public void setEnd() {
        isEnd = true;
    }

    /**
     * 查看该节点是否为叶子节点
     *
     * @return
     */
    public boolean isEnd() {
        return isEnd;
    }
}
